{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 文本相似度 3  高效多文本相似度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 导入相关包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer # 导入trainer相关的包\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from datasets import load_dataset\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载数据集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dir = \"/data/datasets/SimCLUE/datasets/train_pair_1w.json\"\n",
    "dataset = load_dataset('json', data_files=dataset_dir,split='train')\n",
    "print(dataset)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 划分数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = dataset.train_test_split(test_size=0.2)\n",
    "datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建dataloader\n",
    "import torch\n",
    "tokenizer = AutoTokenizer.from_pretrained('/data/models/huggingface/chinese-macbert-base')\n",
    "def process_function(examples):\n",
    "    sententces = []\n",
    "    labels = []\n",
    "    for sen1,sen2,label in zip(examples['sentence1'],examples['sentence2'],examples['label']):\n",
    "        sententces.append(sen1)\n",
    "        sententces.append(sen2)\n",
    "        labels.append(1 if int(label) == 1 else -1)\n",
    "    \n",
    "    tokenizer_examples = tokenizer(sententces, max_length = 128 , truncation=True, padding=\"max_length\")\n",
    "    tokenizer_examples =  {k : [v[i:i+2] for i in range(0,len(v),2)] for k,v in tokenizer_examples.items()}\n",
    "    tokenizer_examples[\"labels\"] = labels\n",
    "    return tokenizer_examples\n",
    "\n",
    "tokenizer_datasets = datasets.map(process_function, batched=True,remove_columns=datasets['train'].column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_data = tokenizer_datasets['train'][0]\n",
    "first_data.keys()\n",
    "print(len(first_data['input_ids']))\n",
    "print(len(first_data['token_type_ids']))\n",
    "print(len(first_data['attention_mask']))\n",
    "print(first_data['labels'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建模型 (自定义模型,重点)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertForSequenceClassification,BertPreTrainedModel,BertModel\n",
    "from typing import Optional, Union\n",
    "from torch.nn import CosineSimilarity,CosineEmbeddingLoss\n",
    "\n",
    "model =  AutoModelForSequenceClassification.from_pretrained('/data/models/huggingface/chinese-macbert-base',num_labels=1)\n",
    "class DualModel(BertPreTrainedModel):\n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "        self.num_labels = config.num_labels\n",
    "        self.config = config\n",
    "        self.bert = BertModel(config)\n",
    "        self.post_init()\n",
    "    \n",
    "    def forward(\n",
    "        self,\n",
    "        input_ids: Optional[torch.Tensor] = None,\n",
    "        attention_mask: Optional[torch.Tensor] = None,\n",
    "        token_type_ids: Optional[torch.Tensor] = None,\n",
    "        position_ids: Optional[torch.Tensor] = None,\n",
    "        head_mask: Optional[torch.Tensor] = None,\n",
    "        inputs_embeds: Optional[torch.Tensor] = None,\n",
    "        labels: Optional[torch.Tensor] = None,\n",
    "        output_attentions: Optional[bool] = None,\n",
    "        output_hidden_states: Optional[bool] = None,\n",
    "        return_dict: Optional[bool] = None,\n",
    "    ):\n",
    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
    "\n",
    "        # 分别获取sentence1 和 sentence2 的输入\n",
    "        senA_input_ids,senB_input_ids = input_ids[:,0],input_ids[:,1]\n",
    "        senA_attention_mask,senB_attention_mask = attention_mask[:,0],attention_mask[:,1]\n",
    "        senA_token_type_ids,senB_token_type_ids = token_type_ids[:,0],token_type_ids[:,1]\n",
    "\n",
    "        outputsA = self.bert(\n",
    "            senA_input_ids,\n",
    "            attention_mask=senA_attention_mask,\n",
    "            token_type_ids=senA_token_type_ids,\n",
    "            position_ids=position_ids,\n",
    "            head_mask=head_mask,\n",
    "            inputs_embeds=inputs_embeds,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            return_dict=return_dict,\n",
    "        )\n",
    "        senA_pooled_output = outputsA[1] # [batch,hidden]\n",
    "\n",
    "        outputsB = self.bert(\n",
    "            senB_input_ids,\n",
    "            attention_mask=senB_attention_mask,\n",
    "            token_type_ids=senB_token_type_ids,\n",
    "            position_ids=position_ids,\n",
    "            head_mask=head_mask,\n",
    "            inputs_embeds=inputs_embeds,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            return_dict=return_dict,\n",
    "        )\n",
    "        senB_pooled_output = outputsB[1] # [batch,hidden]\n",
    "\n",
    "\n",
    "        # 计算相似度 使用torch的库\n",
    "        cos =   CosineSimilarity()(senA_pooled_output,senB_pooled_output)  # 两个向量的余弦相似度\n",
    "\n",
    "        # print(senA_pooled_output.shape,senB_pooled_output.shape,labels.shape)\n",
    "        # 计算loss\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss_fct =  CosineEmbeddingLoss(0.3)\n",
    "            # torch.Size([32, 768]) torch.Size([32, 768]) torch.Size([32])\n",
    "            loss =  loss_fct(senA_pooled_output,senB_pooled_output,labels)\n",
    "\n",
    "\n",
    "        output = (cos,)\n",
    "        return ((loss,) + output) if loss is not None else output\n",
    "\n",
    "model =  DualModel.from_pretrained('/data/models/huggingface/chinese-macbert-base')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估部分重写evaluate的部分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "# 4 项目都加入\n",
    "acc_metric =  evaluate.load('/evaluate-main/metrics/accuracy')\n",
    "f1_metric =  evaluate.load('/evaluate-main/metrics/f1')\n",
    "\n",
    "# 写一个多种标准合成的评价函数\n",
    "def eval_metrics(eval_predict):\n",
    "    predictions,labels =  eval_predict\n",
    "    predictions  = [ int(p>0.7) for p in predictions] \n",
    "    # l 大于 0 为 1 小于 0 为 0 由于原始标签是 1和0 ,但是计算CosineSimilarity之后的标签是 -1 1 所以需要转换\n",
    "    labels = [ int(l>0) for l in labels] \n",
    "    \n",
    "    combine_acc = acc_metric.compute(predictions=predictions, references=labels)\n",
    "    f1 = f1_metric.compute(predictions=predictions, references=labels)\n",
    "    combine_acc.update(f1)\n",
    "    return combine_acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练和验证"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建 TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_args =  TrainingArguments(\n",
    "    output_dir='D:/code/logs/dual_model/checkpoints', # 训练模型的输出路径\n",
    "    per_device_train_batch_size=32, # 训练集 batch size\n",
    "    per_device_eval_batch_size=32, # 验证集 batch size\n",
    "    num_train_epochs=5, # 训练的 epochs 数    \n",
    "    logging_steps=150, # 每 150 个 batch 记录一次日志\n",
    "    # save_steps=100, # 每 100 个 batch 保存一次模型 和每个epoch 保存一次模型互斥\n",
    "    save_strategy = 'epoch', # 每 1 个 epoch 保存一次模型\n",
    "    save_total_limit=3, # 保存 1 个模型\n",
    "    evaluation_strategy='epoch', # 每 1 个 epoch 验证一次模型\n",
    "    learning_rate=4e-5, # 学习率\n",
    "    weight_decay=0.01, # 权重衰减\n",
    "    metric_for_best_model='f1',\n",
    "    load_best_model_at_end=True, # 加载最优模型\n",
    "    report_to='tensorboard', # 训练日志保存到 tensorboard\n",
    "    )\n",
    "train_args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建  Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorWithPadding\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model, # 指定模型\n",
    "    args=train_args, # 执行训练参数\n",
    "    train_dataset=tokenizer_datasets['train'], # 训练数据集指定 切分后的数据集\n",
    "    eval_dataset=tokenizer_datasets['test'], # 验证数据集指定\n",
    "    compute_metrics = eval_metrics, # tokenizer\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate()\n",
    "trainer.evaluate(tokenizer_datasets['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import CosineSimilarity\n",
    "class SentenceSimilarityPipeline:\n",
    "    def __init__(self, model, tokenizer, device):\n",
    "        self.model = model.bert\n",
    "        self.tokenizer = tokenizer\n",
    "        self.device = device\n",
    "\n",
    "    def pre_process(self,senA, SenB):\n",
    "        return self.tokenizer([senA, SenB], max_length = 128 ,padding=True,return_tensors=\"pt\").to(self.device)\n",
    "\n",
    "    def predict(self, inputs):\n",
    "        inputs = {k:v.to(self.device) for k,v in inputs.items()}\n",
    "        return self.model(**inputs)[1] # [2,768]\n",
    "\n",
    "    def post_process(self, logits):\n",
    "        cos =  CosineSimilarity()(logits[None,0,:],logits[None,1,:]).squeeze().cpu().item()\n",
    "        return cos \n",
    "\n",
    "    def __call__(self, senA, SenB,return_vector=False):\n",
    "        inputs = self.pre_process(senA, SenB)\n",
    "        outputs = self.predict(inputs)\n",
    "        result = self.post_process(outputs)\n",
    "        if return_vector:\n",
    "            return result,outputs\n",
    "        else:\n",
    "            return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "model.config.id2label = {0: '不相似', 1: '相似'}\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "pipe = SentenceSimilarityPipeline(model,tokenizer,device)\n",
    "res  = pipe(\n",
    "    \"我喜欢北京，这不仅仅是因为我是一个北京人，住在北京.而是因为北京有着悠久的历史和人文景观，还有它的生活节奏和生活氛围我都喜欢。\",\n",
    "    \"我是一个北京人，所以我喜欢这里。\",return_vector=True)\n",
    "print(res[1].shape)\n",
    "# res[1] 代表了向量的表示"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
